#include <Ntifs.h>
#include <ntddk.h>
#include <ntddkbd.h>
#include <ntddmou.h>
#include <ntdd8042.h>

#include <stdarg.h>
#include <stdio.h>

#include "nmkb.h"

#ifdef NDEBUG
#define DPrint(x,...) 
#else
#define DPrint(x,...) DbgPrintEx(DPFLTR_IHVVIDEO_ID,DPFLTR_ERROR_LEVEL ,x"\n",__VA_ARGS__)
#endif

#ifdef _X86_
extern "C" ULONG __security_cookie=__LINE__;
#endif

static PDEVICE_OBJECT i8prt_dev=0;
static PDEVICE_OBJECT kbd_dev = 0;
static PVOID *kbd_htarget = 0;
static PVOID class_service = 0;
static int installed = 0;
static PDRIVER_OBJECT video_drv = 0;
static PDEVICE_OBJECT video0 = 0,video_pdo0 = 0, video_uname0;
static PDEVICE_OBJECT monitor0 = 0;

static PDRIVER_DISPATCH DpiGdoDispatchInternalIoctl_orig = 0;
static PDRIVER_DISPATCH *DpiGdoDispatchInternalIoctl_loc = 0;

static wchar_t pointer[] = L"\\Device\\PointerClass0";
static unsigned offset = 0x58;

static UNICODE_STRING ukeybd,udriver,upointer;

static int hook_mouse = 0;
RTL_QUERY_REGISTRY_TABLE qtable[]={
	{0,RTL_QUERY_REGISTRY_DIRECT,L"mouse",&upointer,REG_SZ,&pointer,0},
	{0,RTL_QUERY_REGISTRY_DIRECT,L"hookmouse",&hook_mouse,REG_DWORD,0,0},
	{0,RTL_QUERY_REGISTRY_DIRECT,L"hidoffset",&offset,REG_DWORD,(void*)0x58,0},
	{0}
};
struct on_off_t{
    bool value;
    void toggle(){value = !value;}
    bool on(){return value;}
    bool set_off(){return value=false;}
    bool set_on(){return value = true;}
};
static on_off_t r_on_off = {true};
static on_off_t l_on_off = {true};

// BasicDisplay!BASIC_DISPLAY_DRIVER::BlackOutScreen+0xfa (FPO: [Non-Fpo])
// BasicDisplay!BASIC_DISPLAY_DRIVER::SetVidPnSourceVisibility+0x8d (FPO: [1,2,0])
// BasicDisplay!BddDdiSetVidPnSourceVisibility+0x43 (FPO: [Non-Fpo])
// dxgkrnl!ADAPTER_DISPLAY::DdiSetVidPnSourceVisibility+0x44 (FPO: [Non-Fpo])
// dxgkrnl!VIDPN_MGR::RemovePresentPathFromActiveVidPn+0x13f (FPO: [Non-Fpo])
// dxgkrnl!VIDPN_MGR::_MonitorEventHandler+0xd5 (FPO: [6,9,4])
// dxgkrnl!MONITOR_MGR::_IssueMonitorEvent+0x75 (FPO: [Non-Fpo])
// dxgkrnl!MONITOR_MGR::_EnableDisableMonitorInternal+0xb9 (FPO: [Non-Fpo])
// dxgkrnl!MONITOR_MGR::_EnableDisableMonitor+0xe9 (FPO: [1,10,0])
// dxgkrnl!MonitorEnableDisableMonitor+0x64 (FPO: [Non-Fpo])
// dxgkrnl!DpiGdoDispatchInternalIoctl+0x3b7 (FPO: [Non-Fpo])
// dxgkrnl!DpiDispatchInternalIoctl+0x7c (FPO: [2,41,4])
// win32k!GreDeviceIoControlEx+0x61 (FPO: [Non-Fpo])
// win32k!DrvSetMonitorPowerState+0xab (FPO: [Non-Fpo])
// win32k!PowerOffMonitor+0xc4 (FPO: [Non-Fpo])
// win32k!xxxUserPowerEventCalloutWorker+0x2378a5
// win32k!xxxUserPowerCalloutWorker+0x75 (FPO: [0,7,4])
// win32k!NtUserCallNoParam+0x1ab (FPO: [Non-Fpo])
// nt!KiFastCallEntry+0x12c (FPO: [0,3] TrapFrame @ 8cba2d54)
// ntdll!KiFastSystemCallRet (FPO: [0,0,0])
// winsrv!NtUserCallNoParam+0xa (FPO: [1,0,0])
// winsrv!NotificationThread+0x29e (FPO: [1,163,4])
// ntdll!__RtlUserThreadStart+0x51 (FPO: [SEH])
// ntdll!_RtlUserThreadStart+0x1c (FPO: [Non-Fpo])


NTSTATUS  DpiGdoDispatchInternalIoctl_my(DEVICE_OBJECT *DeviceObject,IRP *Irp)
{
    PIO_STACK_LOCATION sp = IoGetCurrentIrpStackLocation(Irp);
    auto& ioctl  = sp->Parameters.DeviceIoControl;
    if (sp->MajorFunction == IRP_MJ_INTERNAL_DEVICE_CONTROL && video_drv == DeviceObject->DriverObject
        && ioctl.IoControlCode == 0x23200F && ioctl.InputBufferLength == 4)
    {
        auto* ps = (POWER_STATE*)ioctl.Type3InputBuffer;
        if (ps)
        {
            if (ps->DeviceState == PowerDeviceD0)
            {
                DPrint("Power request intercepted.: %Ix",DeviceObject);
                Irp->IoStatus.Status = STATUS_SUCCESS;
                Irp->IoStatus.Information = 0;
                IoCompleteRequest(Irp,IO_NO_INCREMENT);
                return STATUS_SUCCESS;
            }
            else if (ps->DeviceState == PowerDeviceD3)
            {
                r_on_off.set_off();
            }
            else if (ps->DeviceState & 0xf0)
            {
                (ULONG&)ps->DeviceState &=  0xf;
            }
            DPrint("REAL request: %Ix,%Ix",*(PULONG)ioctl.Type3InputBuffer,DeviceObject);
        }
    }
    if (DpiGdoDispatchInternalIoctl_orig)
    {
        return DpiGdoDispatchInternalIoctl_orig(DeviceObject,Irp);
    }
    else
    {
        Irp->IoStatus.Status = STATUS_SUCCESS;
        Irp->IoStatus.Information = 0;
        IoCompleteRequest(Irp,IO_NO_INCREMENT);
        return STATUS_SUCCESS;
    }
}

void hook_DpiGdoDispatchInternalIoctl(PDRIVER_OBJECT d)
{
    if (!d || DpiGdoDispatchInternalIoctl_loc) return;
    __try{
        auto drv = d;
        auto p =  drv->MajorFunction[IRP_MJ_INTERNAL_DEVICE_CONTROL];
        if (p)
        {
            DpiGdoDispatchInternalIoctl_orig = p;
            DpiGdoDispatchInternalIoctl_loc = & (drv->MajorFunction[IRP_MJ_INTERNAL_DEVICE_CONTROL]);
            drv->MajorFunction[IRP_MJ_INTERNAL_DEVICE_CONTROL] = DpiGdoDispatchInternalIoctl_my;
            DPrint("Hooked DpiGdoDispatchInternalIoctl:%Ix",DpiGdoDispatchInternalIoctl_loc);
        }
    }
    __except(EXCEPTION_EXECUTE_HANDLER){
        DPrint("Hook DpiGdoDispatchInternalIoctl failed.");
    }
}

VOID rset_sus( )
{
	DPrint("In set_sus IRQL: %d",KeGetCurrentIrql());
	if (video_drv == 0)
	{
		DPrint("video_drv not ready");
		return;
	}
    DPrint("Status: %Ix,%Ix",l_on_off.value,r_on_off.value);
    l_on_off.value = r_on_off.value;
    DEVICE_POWER_STATE state;
    if( l_on_off.on() )
    {
        state = PowerDeviceD3;
    }
    else {
        state = PowerDeviceD0;
    }
	for (auto p = video_drv->DeviceObject; p != 0;p = p->NextDevice)
	{

		if (p->DeviceType != FILE_DEVICE_VIDEO)
			continue;
		auto target = p;

		IO_STATUS_BLOCK ios;
		ULONG aa = 0;
		KEVENT event;
		KeInitializeEvent(&event, NotificationEvent, FALSE);
        POWER_STATE ps;
        ps.DeviceState = (DEVICE_POWER_STATE) (state | 0xf0);
		PIRP irp = IoBuildDeviceIoControlRequest(0x23200F,target,&ps,4,&aa,1,TRUE,&event,&ios);
		if (irp){
            DPrint("Before IoCallDriver State: %Ix",ps.DeviceState);
			NTSTATUS status = IoCallDriver(target,irp);
			if (status == STATUS_PENDING) {
				KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
				status = ios.Status;
			}
			if (status == STATUS_SUCCESS)
            {
                DPrint("Send IRP for %Ix",target);
            }
			else 
				DPrint("Failed sent	IRP:%Ix,%Ix",target,status);
		}
	}
    l_on_off.toggle();
    r_on_off.value = l_on_off.value;
    DPrint("Status: %Ix,%Ix",l_on_off.value,r_on_off.value);
}

VOID set_sus()
{
	PIO_WORKITEM item = IoAllocateWorkItem(kbd_dev);
	if (item)
	{
		IoQueueWorkItem(item,[](PDEVICE_OBJECT,PVOID Context)
		{
			rset_sus();
			IoFreeWorkItem((PIO_WORKITEM)Context);
		}
		,DelayedWorkQueue,item);
	}
}

static NTSTATUS rset_sus2_comp(PDEVICE_OBJECT DeviceObject,PIRP irp,PVOID  	Context )
{
	IoFreeIrp(irp);
	DPrint("Power complete");
	return STATUS_SUCCESS;
}
VOID rset_sus2()
{
	if (video_pdo0 == 0)
	{
		DPrint("videopdo0 is not ready");
	}
	PDEVICE_OBJECT target = video_pdo0;
	POWER_STATE ps;
	ps.DeviceState = PowerDeviceD3;
	//PIRP irp = IoBuildAsynchronousFsdRequest(IRP_MJ_POWER,target,NULL,0,NULL,NULL);
	PIRP irp =IoAllocateIrp(target->StackSize,0);
	if (irp)
	{
		IoSetCompletionRoutine(irp, rset_sus2_comp, rset_sus2_comp, TRUE, TRUE, TRUE);
		irp->IoStatus.Status = STATUS_NOT_IMPLEMENTED;
		irp->IoStatus.Information = 0;

		irp->Flags = IRP_SYNCHRONOUS_API | IRP_BUFFERED_IO;
		irp->RequestorMode = KernelMode;
		irp->UserIosb = 0;
		irp->UserEvent = 0;
		irp->Tail.Overlay.OriginalFileObject = 0;
		irp->Tail.Overlay.Thread = PsGetCurrentThread();
		irp->AssociatedIrp.SystemBuffer = 0;

		PIO_STACK_LOCATION Stack = IoGetNextIrpStackLocation(irp);
		Stack->MajorFunction = IRP_MJ_POWER;
		Stack->DeviceObject = target;
		Stack->MinorFunction = IRP_MN_SET_POWER;

		Stack->Parameters.Power.Type = DevicePowerState;
		Stack->Parameters.Power.State = ps;
		Stack->Parameters.Power.ShutdownType = PowerActionNone;
		Stack->Parameters.Power.SystemContext = 0;
		PoCallDriver(target, irp);
		DPrint("Sent power irp");
	}
}

VOID OnUnload(PDRIVER_OBJECT DriverObject)
{
	if (kbd_htarget)
	{
		*kbd_htarget = class_service;
	}
    if (DpiGdoDispatchInternalIoctl_loc && DpiGdoDispatchInternalIoctl_orig)
    {
        *DpiGdoDispatchInternalIoctl_loc = DpiGdoDispatchInternalIoctl_orig;
    }
}

PDEVICE_OBJECT find_root(PDEVICE_OBJECT d)
{
	while (d)
	{
		auto t = d->DeviceObjectExtension->AttachedTo;
		if (!t) return d;
		d = t;
	}
	return d;
}

// VOID find_video()
// {
// 	NTSTATUS Status;
// 	PDEVICE_OBJECT p = 0;
// 	if (video_drv == 0) return;
// 
// 	p = video_drv->DeviceObject;
// 	while (p)
// 	{
// 		ULONG l = 0;
// 		CHAR nameBuffer[256+sizeof( OBJECT_NAME_INFORMATION )] = {0};
// 		POBJECT_NAME_INFORMATION pni = (POBJECT_NAME_INFORMATION) &nameBuffer[0];
// 		Status = ObQueryNameString(p,pni,sizeof(nameBuffer),&l);
// 		if (Status == STATUS_SUCCESS && RtlEqualUnicodeString(&umonitor,&pni->Name,TRUE)) /*videopdo1*/
// 		{
// 			video0 = p;
// 			DPrint("Find video0:%Ix",video0);
// 		}
// 		else if (Status == 0 &&  pni->Name.Length == 0)
// 		{
// 			video_uname0 = p;
// 			DPrint ("Find unamed video:%Ix",video_uname0);
// 		}
// 		p = p->NextDevice;
// 	}
// }

VOID find_monitor()
{
	UNICODE_STRING monitor_drv_name;
	PDRIVER_OBJECT monitor_drv = 0;
	RtlInitUnicodeString(&monitor_drv_name,L"\\Driver\\monitor");
	NTSTATUS stat = 0;

	stat = ObReferenceObjectByName(&monitor_drv_name,OBJ_CASE_INSENSITIVE,0,0,*IoDriverObjectType,KernelMode,0,(PVOID*)&monitor_drv);

	if (stat != STATUS_SUCCESS || monitor_drv == 0) 
	{
		DPrint("Failed to find monitor driver (%Ix)",stat);
		return;
	}

	PDEVICE_OBJECT p = monitor_drv->DeviceObject;

	while (p)
	{
		auto root = find_root(p);
		if (root && root != p && root->DeviceType == FILE_DEVICE_SCREEN) 
		{
			monitor0 = p;
			DPrint("Found monitor %Ix",monitor0);
			video_drv = root->DriverObject;
			DPrint("Found video driver: %Ix",video_drv);
            hook_DpiGdoDispatchInternalIoctl(video_drv);
			video_pdo0 = root;
			DPrint("Found video_pdo:%Ix",video_pdo0);
			goto end;
		}
		p = p->NextDevice;
	}
end:
	ObDereferenceObject(monitor_drv);
}

VOID find_kbd()
{
	UNICODE_STRING i8prt_drv_name , kbd_drv_name;
	PDRIVER_OBJECT kbd_drv = 0 , i8prt_drv = 0;
	RtlInitUnicodeString(&i8prt_drv_name,L"\\Driver\\i8042prt");
	RtlInitUnicodeString(&kbd_drv_name,L"\\Driver\\kbdclass");
	NTSTATUS stat = 0;

	stat = ObReferenceObjectByName(&i8prt_drv_name,OBJ_CASE_INSENSITIVE,0,0,*IoDriverObjectType,KernelMode,0,(PVOID*)&i8prt_drv);
	if (stat != STATUS_SUCCESS || i8prt_drv == 0)
	{
		DPrint("Failed to get i8042prt driver:%Ix",stat);
		return;
	}

	stat = ObReferenceObjectByName(&kbd_drv_name,OBJ_CASE_INSENSITIVE,0,0,*IoDriverObjectType,KernelMode,0,(PVOID*)&kbd_drv);
	if (stat != STATUS_SUCCESS || kbd_drv == 0)
	{
		DPrint("Failed to get kbdclass driver:%Ix",stat);
		return;
	}

	PDEVICE_OBJECT p = i8prt_drv->DeviceObject;

	while (p)
	{
		PDEVICE_OBJECT u = p->AttachedDevice;
		while (u)
		{
			PDRIVER_OBJECT drv = u->DriverObject;
			if (drv == kbd_drv)
			{
				i8prt_dev = p;
				kbd_dev = u;
				goto end;
			}
			u = u->AttachedDevice;
		}
		p = p->NextDevice;
	}
end:
	ObDereferenceObject(kbd_drv);
	ObDereferenceObject(i8prt_drv);
}


VOID KeyboardClassServiceCallback(PDEVICE_OBJECT dev,PKEYBOARD_INPUT_DATA InputDataStart,PKEYBOARD_INPUT_DATA InputDataEnd,PULONG InputDataConsumed)
{	
	((void(*)(PDEVICE_OBJECT,PKEYBOARD_INPUT_DATA,PKEYBOARD_INPUT_DATA,PULONG))class_service)(dev,InputDataStart,InputDataEnd,InputDataConsumed);
	static int htcount = 0;
	static  ULONGLONG tick = 0;

	ULONG count = * InputDataConsumed;
	for (int i = 0; i < count; i++)
	{
		auto scan = InputDataStart[i].MakeCode;
		auto flags = InputDataStart[i].Flags;
		if(scan == 0x1d && (flags == (KEY_E0|KEY_BREAK))) {
			ULONGLONG curr = KeQueryInterruptTime();
			if(curr-tick < 5000000i64){
				htcount++;
			}
			else {
				htcount = 0;
			}

			tick = curr;
			if(htcount >= 2) {
				htcount = 0;
				set_sus();
			}
		}
	}
}
VOID hook_kbd()
{
	PULONG pd = (PULONG)i8prt_dev->DeviceExtension;
	for (int i = 0;i < 0x100; i++)
	{
		PVOID * pp = (PVOID*) pd;
		PVOID d1 = *pp;
		if (d1 == kbd_dev)
		{
			kbd_htarget = pp + 1;
			break;
		}
		pd++;
	}
	if (kbd_htarget)
	{
		DPrint("Found hook point:%Ix",kbd_htarget);
		class_service = *kbd_htarget;
		*kbd_htarget = KeyboardClassServiceCallback;
	}
}

NTSTATUS NTAPI DriverEntry(PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegistryPath )
{
	DPrint("nmkb2 DriverEntry");

	NTSTATUS ret = RtlQueryRegistryValues(RTL_REGISTRY_ABSOLUTE,RegistryPath->Buffer,qtable,0,0);
	if(ret != STATUS_SUCCESS) {
		DPrint("Registry query failed: %S %Ix",RegistryPath->Buffer,ret);
		return ret;
	}
	
	find_monitor();
// 	find_video();

 	if (video_drv == 0)
	{
		DPrint("video_drv not found.");
		return STATUS_DEVICE_NOT_CONNECTED;
	}

	//rset_sus();
	//return STATUS_DEVICE_NOT_CONNECTED;

	find_kbd();
	if(kbd_dev == 0 || i8prt_dev == 0) {
		DPrint("keyboard not found");
		return STATUS_DEVICE_NOT_CONNECTED;
	}

	hook_kbd();

	DriverObject->DriverUnload = OnUnload;

	return STATUS_SUCCESS;
}